import os
from ultralytics import YOLO

# 定义权重文件路径和导出路径
weight_path = r'./best.pt'  # 替换为您的权重文件路径
onnx_output_path = 'best.onnx'  # 输出的 ONNX 文件名

# 确保权重文件存在
if not os.path.exists(weight_path):
    print(f"权重文件 {weight_path} 不存在。")
else:
    try:
        # 加载模型
        model = YOLO(weight_path)

        # 导出为 ONNX 格式，使用适当的参数
        model.export(format='onnx', imgsz=(640, 640), dynamic=True)

        print(f"权重文件成功转换为 ONNX 格式，输出路径: {onnx_output_path}")

    except Exception as e:
        print(f"转换过程中出现错误: {e}")
